from sklearn.metrics import f1_score, accuracy_score, cohen_kappa_score, confusion_matrix

def calculate_metrics(y_true, y_pred, classes=['W', 'N1', 'N2', 'N3', 'REM']):
    """计算所有评估指标"""
    # 总体指标
    overall_metrics = {
        'f1_score': f1_score(y_true, y_pred, average='weighted'),
        'accuracy': accuracy_score(y_true, y_pred),
        'kappa': cohen_kappa_score(y_true, y_pred)
    }
    
    # 每个类别的F1分数
    per_class_f1 = f1_score(y_true, y_pred, average=None)
    per_class_metrics = {cls: per_class_f1[i] for i, cls in enumerate(classes)}
    
    # 混淆矩阵
    cm = confusion_matrix(y_true, y_pred)
    
    return {
        'overall': overall_metrics,
        'per_class': per_class_metrics,
        'confusion_matrix': cm
    }
